#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from std_msgs.msg import Bool
from geometry_msgs.msg import Pose
from moveit_commander.move_group import MoveGroupCommander
import moveit_commander
import sys
import numpy as np
from tf_transformations import quaternion_matrix
from rclpy.callback_groups import ReentrantCallbackGroup
from rclpy.executors import MultiThreadedExecutor

class MoveArmNode(Node):
    def __init__(self):
        super().__init__('move_arm_to_poses')
        
        # MoveIt 初始化
        moveit_commander.roscpp_initialize(sys.argv)
        self.move_group = MoveGroupCommander("right_arm")
        self.move_group.set_end_effector_link("right_tool_link")
        self.move_group.set_planning_time(0.1)
        self.move_group.set_max_velocity_scaling_factor(0.8)
        self.move_group.set_max_acceleration_scaling_factor(0.8)
        
        # 使用回调组处理并发回调
        self.callback_group = ReentrantCallbackGroup()
        
        # 创建发布者
        self.hand_control_pub = self.create_publisher(
            Bool, 
            "/hand_control_topic", 
            10
        )
        
        # 创建订阅者
        self.right_pose_sub = self.create_subscription(
            Pose,
            '/right_grasp_pose',
            self.right_pose_callback,
            10,
            callback_group=self.callback_group
        )
        
        self.task = "grasp"
        self.get_logger().info("Node initialized")

    def right_pose_callback(self, msg):
        self.get_logger().info("Received new right pose")
        quaternion = [
            msg.orientation.x,
            msg.orientation.y,
            msg.orientation.z,
            msg.orientation.w
        ]
        
        rotation_matrix = quaternion_matrix(quaternion)
        local_offset = np.array([0.15, 0.0, 0.0, 1.0])
        world_offset = np.dot(rotation_matrix, local_offset)

        if self.task == "grasp":
            self._handle_grasp_task(msg, world_offset)
        elif self.task == "release":
            self._handle_release_task()

    def _handle_grasp_task(self, msg, offset):
        try:
            # 控制夹爪闭合
            hand_msg = Bool()
            hand_msg.data = True
            self.hand_control_pub.publish(hand_msg)
            
            # 计算目标位姿
            target_pose = Pose()
            target_pose.position.x = msg.position.x + offset[0]
            target_pose.position.y = msg.position.y + offset[1]
            target_pose.position.z = msg.position.z + offset[2]
            target_pose.orientation = msg.orientation
            
            # 执行运动
            if self.move_to_pose(target_pose):
                self.get_logger().info("Grasp movement completed")
                self.task = "release"
        except Exception as e:
            self.get_logger().error(f"Grasp task failed: {str(e)}")

    def _handle_release_task(self):
        try:
            # 控制夹爪张开
            hand_msg = Bool()
            hand_msg.data = False
            self.hand_control_pub.publish(hand_msg)
            
            # 此处可添加释放物体的运动逻辑
            self.get_logger().info("Release movement completed")
            self.task = "grasp"
        except Exception as e:
            self.get_logger().error(f"Release task failed: {str(e)}")

    def move_to_pose(self, pose):
        """执行运动到指定位姿"""
        self.move_group.set_pose_target(pose)
        
        # 运动规划
        plan_result = self.move_group.plan()
        
        if plan_result[0]:
            self.get_logger().info(f"Planning succeeded in {plan_result[2]}s")
            move_result = self.move_group.execute(plan_result[1], wait=True)
            
            if move_result:
                self.get_logger().info("Execution completed successfully")
                return True
            else:
                self.get_logger().warn("Execution failed")
                return False
        else:
            self.get_logger().error("Planning failed")
            return False

def main(args=None):
    rclpy.init(args=args)
    
    try:
        # 使用多线程执行器处理并发回调
        executor = MultiThreadedExecutor()
        node = MoveArmNode()
        executor.add_node(node)
        
        try:
            executor.spin()
        finally:
            executor.shutdown()
            node.destroy_node()
            
    except KeyboardInterrupt:
        pass
    finally:
        rclpy.shutdown()

if __name__ == '__main__':
    main()
